#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""
Enhanced P4Runtime controller for advanced_tunnel.p4.
- Supports 3 switches (s1, s2, s3) each with one directly attached host (port 1).
- Programs tunnels for ALL 6 host-to-host directions (h1<->h2, h1<->h3, h2<->h3).
- Uses distinct tunnel IDs per *directed* link:
    s1->s2:100, s2->s1:101
    s1->s3:200, s3->s1:201
    s2->s3:300, s3->s2:301
- Reads ingress/egress counters for every direction and appends results
  to per-link log files: S1S2.txt, S1S3.txt, S2S3.txt (both directions).
  (One record per 2 seconds; safe to tail -f)
NOTE: Port map assumes Mininet triangle topology from P4 tutorials:
      host on port 1; s1<->s2 on port 2; s1<->s3 on port 3; s2<->s3 on port 3,
      and symmetric at the peer side (common in the tutorial).
      If your topology differs, adjust PORT_MAP below.
"""
import argparse
import os
import sys
from time import sleep
from datetime import datetime

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# ---------- Topology assumptions (edit if needed) ----------
PORT_MAP = {
    's1': {'h': 1, 's2': 2, 's3': 3},
    's2': {'h': 1, 's1': 2, 's3': 3},
    's3': {'h': 1, 's1': 2, 's2': 3},
}

# Hosts (standard tutorial addresses; edit if needed)
HOSTS = {
    'h1': {'ip': '10.0.1.1', 'mac': '08:00:00:00:01:11', 'sw': 's1'},
    'h2': {'ip': '10.0.2.2', 'mac': '08:00:00:00:02:22', 'sw': 's2'},
    'h3': {'ip': '10.0.3.3', 'mac': '08:00:00:00:03:33', 'sw': 's3'},
}

# Directed tunnel IDs per *switch->switch* direction
TUNNEL_ID = {
    ('s1','s2'): 100, ('s2','s1'): 101,
    ('s1','s3'): 200, ('s3','s1'): 201,
    ('s2','s3'): 300, ('s3','s2'): 301,
}

LINK_FILES = {
    frozenset(('s1','s2')): 'S1S2.txt',
    frozenset(('s1','s3')): 'S1S3.txt',
    frozenset(('s2','s3')): 'S2S3.txt',
}


def add_ingress_rule(p4info_helper, sw_conn, dst_ip, tunnel_id):
    """ipv4_lpm -> myTunnel_ingress(dst_id)"""
    entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip, 32)},
        action_name="MyIngress.myTunnel_ingress",
        action_params={"dst_id": tunnel_id}
    )
    sw_conn.WriteTableEntry(entry)


def add_transit_rule(p4info_helper, sw_conn, tunnel_id, egress_port):
    """myTunnel_exact(dst_id) -> myTunnel_forward(port)"""
    entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={"hdr.myTunnel.dst_id": tunnel_id},
        action_name="MyIngress.myTunnel_forward",
        action_params={"port": egress_port}
    )
    sw_conn.WriteTableEntry(entry)


def add_egress_rule(p4info_helper, sw_conn, tunnel_id, dst_mac, host_port):
    """myTunnel_exact(dst_id) -> myTunnel_egress(dstAddr, port)"""
    entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={"hdr.myTunnel.dst_id": tunnel_id},
        action_name="MyIngress.myTunnel_egress",
        action_params={"dstAddr": dst_mac, "port": host_port}
    )
    sw_conn.WriteTableEntry(entry)


def program_path(p4info_helper, sw_ing, sw_out, tunnel_id, next_hop_port, dst_ip, dst_mac):
    """Install ingress (ipv4_lpm), transit (forward), egress (egress)"""
    add_ingress_rule(p4info_helper, sw_ing, dst_ip, tunnel_id)
    add_transit_rule(p4info_helper, sw_ing, tunnel_id, next_hop_port)
    add_egress_rule(p4info_helper, sw_out, tunnel_id, dst_mac, PORT_MAP[sw_out.name]['h'])


def read_counter_pair(p4info_helper, a, b, tid):
    """Return (a.ingress[tid], b.egress[tid]) as (packets, bytes) pairs"""
    def _read(sw, cname, idx):
        total_pkts = 0
        total_bytes = 0
        for response in sw.ReadCounters(p4info_helper.get_counters_id(cname), idx):
            for entity in response.entities:
                data = entity.counter_entry.data
                total_pkts += data.packet_count
                total_bytes += data.byte_count
        return total_pkts, total_bytes

    a_in = _read(a, "MyIngress.ingressTunnelCounter", tid)
    b_out = _read(b, "MyIngress.egressTunnelCounter", tid)
    return a_in, b_out


def append_link_log(link_name, msg):
    os.makedirs('p4runtime', exist_ok=True)
    path = os.path.join('p4runtime', f'{link_name}.txt')
    with open(path, 'a', encoding='utf-8') as f:
        f.write(msg + '\n')


def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    # 3 switches
    s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
        name='s1', address='127.0.0.1:50051', device_id=0,
        proto_dump_file='logs/s1-p4runtime-requests.txt')
    s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
        name='s2', address='127.0.0.1:50052', device_id=1,
        proto_dump_file='logs/s2-p4runtime-requests.txt')
    s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
        name='s3', address='127.0.0.1:50053', device_id=2,
        proto_dump_file='logs/s3-p4runtime-requests.txt')

    for sw in (s1, s2, s3):
        sw.MasterArbitrationUpdate()
        sw.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print(f"Installed P4 on {sw.name}")

    # Convenience dict
    sd = {'s1': s1, 's2': s2, 's3': s3}

    # ---------- Program all 6 directions ----------
    # h1->h2: s1 ingress -> s2 egress
    program_path(p4info_helper, sd['s1'], sd['s2'],
                 TUNNEL_ID[('s1','s2')],
                 PORT_MAP['s1']['s2'],
                 HOSTS['h2']['ip'], HOSTS['h2']['mac'])

    # h2->h1
    program_path(p4info_helper, sd['s2'], sd['s1'],
                 TUNNEL_ID[('s2','s1')],
                 PORT_MAP['s2']['s1'],
                 HOSTS['h1']['ip'], HOSTS['h1']['mac'])

    # h1->h3
    program_path(p4info_helper, sd['s1'], sd['s3'],
                 TUNNEL_ID[('s1','s3')],
                 PORT_MAP['s1']['s3'],
                 HOSTS['h3']['ip'], HOSTS['h3']['mac'])

    # h3->h1
    program_path(p4info_helper, sd['s3'], sd['s1'],
                 TUNNEL_ID[('s3','s1')],
                 PORT_MAP['s3']['s1'],
                 HOSTS['h1']['ip'], HOSTS['h1']['mac'])

    # h2->h3
    program_path(p4info_helper, sd['s2'], sd['s3'],
                 TUNNEL_ID[('s2','s3')],
                 PORT_MAP['s2']['s3'],
                 HOSTS['h3']['ip'], HOSTS['h3']['mac'])

    # h3->h2
    program_path(p4info_helper, sd['s3'], sd['s2'],
                 TUNNEL_ID[('s3','s2')],
                 PORT_MAP['s3']['s2'],
                 HOSTS['h2']['ip'], HOSTS['h2']['mac'])

    print("All tunnel rules installed for h1<->h2, h1<->h3, h2<->h3.")

    # ---------- Periodically read & log counters ----------
    try:
        while True:
            ts = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print('\n----- Reading tunnel counters -----')
            # s1<->s2
            a,b = read_counter_pair(p4info_helper, s1, s2, TUNNEL_ID[('s1','s2')])
            c,d = read_counter_pair(p4info_helper, s2, s1, TUNNEL_ID[('s2','s1')])
            append_link_log('S1S2', f'[{ts}] 方向 s1→s2（隧道ID=100）：封装侧(ingress) {a[0]} 包 / {a[1]} 字节；解封侧(egress) {b[0]} 包 / {b[1]} 字节')
            append_link_log('S1S2', f'[{ts}] 方向 s2→s1（隧道ID=101）：封装侧(ingress) {c[0]} 包 / {c[1]} 字节；解封侧(egress) {d[0]} 包 / {d[1]} 字节')

            # s1<->s3
            a,b = read_counter_pair(p4info_helper, s1, s3, TUNNEL_ID[('s1','s3')])
            c,d = read_counter_pair(p4info_helper, s3, s1, TUNNEL_ID[('s3','s1')])
            append_link_log('S1S3', f'[{ts}] 方向 s1→s3（隧道ID=200）：封装侧(ingress) {a[0]} 包 / {a[1]} 字节；解封侧(egress) {b[0]} 包 / {b[1]} 字节')
            append_link_log('S1S3', f'[{ts}] 方向 s3→s1（隧道ID=201）：封装侧(ingress) {c[0]} 包 / {c[1]} 字节；解封侧(egress) {d[0]} 包 / {d[1]} 字节')

            # s2<->s3
            a,b = read_counter_pair(p4info_helper, s2, s3, TUNNEL_ID[('s2','s3')])
            c,d = read_counter_pair(p4info_helper, s3, s2, TUNNEL_ID[('s3','s2')])
            append_link_log('S2S3', f'[{ts}] 方向 s2→s3（隧道ID=300）：封装侧(ingress) {a[0]} 包 / {a[1]} 字节；解封侧(egress) {b[0]} 包 / {b[1]} 字节')
            append_link_log('S2S3', f'[{ts}] 方向 s3→s2（隧道ID=301）：封装侧(ingress) {c[0]} 包 / {c[1]} 字节；解封侧(egress) {d[0]} 包 / {d[1]} 字节')

            # Also print to console for quick inspection
            for link in ('S1S2','S1S3','S2S3'):
                print(f'[{ts}] updated {link}.txt')
            sleep(2)
    except KeyboardInterrupt:
        print("Shutting down.")
    except grpc.RpcError as e:
        print("gRPC Error:", e.details(), end=' ')
        status_code = e.code()
        print("(%s)" % status_code.name, end=' ')
        traceback = sys.exc_info()[2]
        print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))

    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Enhanced P4Runtime Controller')
    parser.add_argument('--p4info', type=str, default='./build/advanced_tunnel.p4.p4info.txtpb',
                        help='p4info proto in text format from p4c')
    parser.add_argument('--bmv2-json', type=str, default='./build/advanced_tunnel.json',
                        help='BMv2 JSON file from p4c')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)

    main(args.p4info, args.bmv2_json)